import numpy as np
from scipy.spatial import cKDTree

class CornerPointGrid:
    def __init__(self, filename):
        self.nx = 0
        self.ny = 0
        self.nz = 0
        self.coords = []  # COORD数据
        self.zcorn = []   # ZCORN数据
        self.actnum = []  # ACTNUM数据
        self.read_grid(filename)

    def read_grid(self, filename):
        """读取角点网格文件"""
        with open(filename, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        reading_section = None
        for line in lines:
            line = line.strip()
            if not line or line == '/':
                reading_section = None
                continue
            
            if line in ['SPECGRID', 'COORD', 'ZCORN', 'ACTNUM']:
                reading_section = line
                continue

            if reading_section == 'SPECGRID':
                dims = line.split()
                self.nx, self.ny, self.nz = map(int, dims[:3])
            
            elif reading_section == 'COORD':
                try:
                    self.coords.extend(map(float, line.split()))
                except ValueError:
                    continue
            
            elif reading_section == 'ZCORN':
                try:
                    self.zcorn.extend(map(float, line.split()))
                except ValueError:
                    continue
            
            elif reading_section == 'ACTNUM':
                for part in line.split():
                    if '*' in part:
                        count, value = part.split('*')
                        self.actnum.extend([int(value)] * int(count))
                    else:
                        try:
                            self.actnum.append(int(part))
                        except ValueError:
                            continue

        self.coords = np.array(self.coords).reshape(-1, 6)
        self.zcorn = np.array(self.zcorn)
        self.actnum = np.array(self.actnum)

    def is_active_cell(self, i, j, k):
        """检查指定单元是否为活动单元"""
        if len(self.actnum) == 0:  # 如果没有ACTNUM数据，认为所有单元都是活动的
            return True
        idx = k * self.nx * self.ny + j * self.nx + i
        return idx < len(self.actnum) and self.actnum[idx] == 1

    def get_cell_vertices(self, i, j, k):
        """获取指定单元的8个顶点坐标
        COMSOL的六面体节点顺序：
        底面（z较小的面）：0,1,2,3 (顺时针)
        顶面（z较大的面）：4,5,6,7 (顺时针)
        """
        vertices = np.zeros((8, 3))
        
        pillar_indices = [
            j * (self.nx + 1) + i,            # 左前
            j * (self.nx + 1) + (i + 1),      # 右前
            (j + 1) * (self.nx + 1) + (i + 1),# 右后
            (j + 1) * (self.nx + 1) + i       # 左后
        ]
        
        nxy = self.nx * self.ny
        base_idx = k * nxy * 8
        top_base_idx = base_idx + nxy * 4
        
        # 修改顶点顺序以匹配COMSOL的期望
        # 底面顶点 (顺时针)
        vertices[0] = [self.coords[pillar_indices[0]][0], self.coords[pillar_indices[0]][1], 
                      self.zcorn[base_idx + j * self.nx * 2 + i * 2]]  # 左前下
        vertices[1] = [self.coords[pillar_indices[1]][0], self.coords[pillar_indices[1]][1], 
                      self.zcorn[base_idx + j * self.nx * 2 + (i+1) * 2]]  # 右前下
        vertices[2] = [self.coords[pillar_indices[2]][0], self.coords[pillar_indices[2]][1], 
                      self.zcorn[base_idx + (j+1) * self.nx * 2 + (i+1) * 2]]  # 右后下
        vertices[3] = [self.coords[pillar_indices[3]][0], self.coords[pillar_indices[3]][1], 
                      self.zcorn[base_idx + (j+1) * self.nx * 2 + i * 2]]  # 左后下
        
        # 顶面顶点 (顺时针)
        vertices[4] = [self.coords[pillar_indices[0]][0], self.coords[pillar_indices[0]][1], 
                      self.zcorn[top_base_idx + j * self.nx * 2 + i * 2]]  # 左前上
        vertices[5] = [self.coords[pillar_indices[1]][0], self.coords[pillar_indices[1]][1], 
                      self.zcorn[top_base_idx + j * self.nx * 2 + (i+1) * 2]]  # 右前上
        vertices[6] = [self.coords[pillar_indices[2]][0], self.coords[pillar_indices[2]][1], 
                      self.zcorn[top_base_idx + (j+1) * self.nx * 2 + (i+1) * 2]]  # 右后上
        vertices[7] = [self.coords[pillar_indices[3]][0], self.coords[pillar_indices[3]][1], 
                      self.zcorn[top_base_idx + (j+1) * self.nx * 2 + i * 2]]  # 左后上

        return vertices

    def calculate_fault_distance(self, fault_surface):
        """计算网格节点到断层面的距离
        
        Args:
            fault_surface: 真实断层面数据
            
        Returns:
            float: 距离的平均值
        """
        distances = []
        nodes, _ = convert_to_fem(self)
        
        for node in nodes:
            # 计算点到面的最短距离
            min_distance = float('inf')
            for triangle in fault_surface:
                dist = point_to_triangle_distance(node, triangle)
                min_distance = min(min_distance, dist)
            distances.append(min_distance)
            
        return np.mean(distances)

def write_vtk_mesh(filename, nodes, elements, fault_markers):
    """输出VTK格式的六面体网格，包含断层节点标记
    
    Args:
        filename: str, 输出文件名
        nodes: np.array, 节点坐标
        elements: np.array, 单元节点连接关系
        fault_markers: np.array, 断层节点标记
    """
    with open(filename, 'w') as f:
        # 1. VTK文件头
        f.write("# vtk DataFile Version 2.0\n")
        f.write("Hexahedral Mesh with Fault Nodes\n")
        f.write("ASCII\n")
        f.write("DATASET UNSTRUCTURED_GRID\n\n")
        
        # 2. 节点坐标
        num_nodes = len(nodes)
        f.write(f"POINTS {num_nodes} float\n")
        for node in nodes:
            f.write(f"{node[0]:.6f} {node[1]:.6f} {node[2]:.6f}\n")
        f.write("\n")
        
        # 3. 单元定义
        num_elements = len(elements)
        total_size = num_elements * 9
        f.write(f"CELLS {num_elements} {total_size}\n")
        for element in elements:
            f.write(f"8 {' '.join(map(str, element))}\n")
        f.write("\n")
        
        # 4. 单元类型
        f.write(f"CELL_TYPES {num_elements}\n")
        for _ in range(num_elements):
            f.write("12\n")
        f.write("\n")
        
        # 5. 添加断层节点标记
        f.write(f"POINT_DATA {num_nodes}\n")
        f.write("SCALARS IsFaultNode int 1\n")
        f.write("LOOKUP_TABLE default\n")
        for marker in fault_markers:
            f.write(f"{int(marker)}\n")

def convert_to_fem(grid):
    """将角点网格转换为六面体有限元网格
    COMSOL六面体单元节点顺序：
    1. 底面逆时针：0-3-2-1
    2. 顶面逆时针：4-7-6-5
    """
    nodes = []
    elements = []
    node_map = {}
    
    # 只处理活动单元
    for k in range(grid.nz):
        for j in range(grid.ny):
            for i in range(grid.nx):
                if not grid.is_active_cell(i, j, k):
                    continue
                    
                # 获取当前单元的8个顶点
                vertices = grid.get_cell_vertices(i, j, k)
                
                # 创建单元，调整节点顺序以匹配COMSOL要求
                comsol_order = [0, 3, 2, 1, 4, 7, 6, 5]  # 重新排序的索引
                element = []
                
                for idx in comsol_order:
                    vertex = vertices[idx]
                    vertex_key = tuple(vertex)
                    if vertex_key not in node_map:
                        node_map[vertex_key] = len(nodes)
                        nodes.append(vertex)
                    element.append(node_map[vertex_key])
                
                elements.append(element)
    
    return np.array(nodes), np.array(elements)

def point_to_triangle_distance(point, triangle):
    """计算点到三角形的最短距离
    
    Args:
        point: np.array, 3D点坐标 [x,y,z]
        triangle: np.array, 三角形三个顶点坐标 [[x1,y1,z1], [x2,y2,z2], [x3,y3,z3]]
        
    Returns:
        float: 点到三角形的最短距离
    """
    # 1. 计算三角形法向量
    v1 = triangle[1] - triangle[0]
    v2 = triangle[2] - triangle[0]
    normal = np.cross(v1, v2)
    normal = normal / np.linalg.norm(normal)
    
    # 2. 计算点到三角形平面的距离
    plane_dist = abs(np.dot(point - triangle[0], normal))
    
    # 3. 计算点在平面上的投影
    proj_point = point - plane_dist * normal
    
    # 4. 判断投影点是否在三角形内
    def is_point_in_triangle(p, triangle):
        """判断点是否在三角形内"""
        v0 = triangle[2] - triangle[0]
        v1 = triangle[1] - triangle[0]
        v2 = p - triangle[0]
        
        dot00 = np.dot(v0, v0)
        dot01 = np.dot(v0, v1)
        dot02 = np.dot(v0, v2)
        dot11 = np.dot(v1, v1)
        dot12 = np.dot(v1, v2)
        
        inv_denom = 1 / (dot00 * dot11 - dot01 * dot01)
        u = (dot11 * dot02 - dot01 * dot12) * inv_denom
        v = (dot00 * dot12 - dot01 * dot02) * inv_denom
        
        return (u >= 0) and (v >= 0) and (u + v <= 1)
    
    if is_point_in_triangle(proj_point, triangle):
        return plane_dist
    
    # 5. 如果不在三角形内,计算到最近边或顶点的距离
    min_dist = float('inf')
    # 检查到三个顶点的距离
    for vertex in triangle:
        dist = np.linalg.norm(point - vertex)
        min_dist = min(min_dist, dist)
    
    # 检查到三条边的距离
    edges = [(triangle[0], triangle[1]),
             (triangle[1], triangle[2]),
             (triangle[2], triangle[0])]
             
    for edge_start, edge_end in edges:
        edge_vec = edge_end - edge_start
        edge_len = np.linalg.norm(edge_vec)
        edge_dir = edge_vec / edge_len
        
        # 算点到边的投影
        t = np.dot(point - edge_start, edge_dir)
        t = max(0, min(t, edge_len))  # 限制在边的范围内
        
        # 计算最近点
        closest = edge_start + t * edge_dir
        dist = np.linalg.norm(point - closest)
        min_dist = min(min_dist, dist)
    
    return min_dist

def get_neighbor_indices(i, j, k, vertex_idx):
    """获取指定顶点的相邻节点索引"""
    neighbors = []
    
    # 根据vertex_idx确定顶点在单元中的位置
    # 返回相邻节点的全局索引
    
    return neighbors

def identify_fault_nodes(grid, nodes):
    """识别断层节点"""
    fault_markers = np.zeros(len(nodes))
    dx = (np.max(grid.coords[:, 0]) - np.min(grid.coords[:, 0])) / grid.nx
    tree = None
    fault_vertices = []
    
    for k in range(grid.nz):
        for j in range(grid.ny-1):
            for i in range(grid.nx-1):
                vertices1 = grid.get_cell_vertices(i, j, k)
                vertices2 = grid.get_cell_vertices(i+1, j, k)
                
                if np.max(np.linalg.norm(vertices2 - vertices1, axis=1)) > 1.2 * dx:
                    fault_vertices.extend(vertices1)
                    
                    if len(fault_vertices) >= 1000:
                        if tree is None:
                            tree = cKDTree(nodes)
                        fault_points = np.array(fault_vertices)
                        distances, indices = tree.query(fault_points, k=1, distance_upper_bound=dx)
                        fault_markers[indices[distances != np.inf]] = 1
                        fault_vertices = []
    
    if fault_vertices:
        if tree is None:
            tree = cKDTree(nodes)
        fault_points = np.array(fault_vertices)
        distances, indices = tree.query(fault_points, k=1, distance_upper_bound=dx)
        fault_markers[indices[distances != np.inf]] = 1
    
    return fault_markers

def main():
    grid = CornerPointGrid("Grid_E.out")
    nodes, elements = convert_to_fem(grid)
    fault_markers = identify_fault_nodes(grid, nodes)
    write_vtk_mesh("finiteGrid.vtk", nodes, elements, fault_markers)

if __name__ == "__main__":
    main()